import os
import re
import shutil
import subprocess
import sys
import torch
import requests
from transformers import AutoModel, AutoTokenizer
from .config_manager import config_manager

class ModelManager:
    """
    模型管理器类，负责模型的加载、检查和下载功能
    """
    def __init__(self, config=None):
        self.config = config or config_manager.config
        self.model = None
        self.tokenizer = None
    
    def check_local_model(self, model_dir=None):
        """
        检查本地是否存在模型，优化版本，支持更多模型文件格式
        """
        # 使用传入的model_dir参数或尝试从配置中获取
        if model_dir is None:
            if 'paths' in self.config and 'model_dir' in self.config['paths']:
                model_dir = self.config['paths']['model_dir']
            elif 'model' in self.config and 'local_model_path' in self.config['model']:
                model_dir = self.config['model']['local_model_path']
            else:
                model_dir = "models"
                print(f"警告: 配置中未找到模型目录，使用默认值: {model_dir}")
                
        print(f"检查本地模型目录: {model_dir}")
        
        # 检查模型目录是否存在
        if not os.path.exists(model_dir):
            print(f"✗ 模型目录不存在")
            return False
        
        # 获取目录中的文件列表
        try:
            files = os.listdir(model_dir)
            print(f"模型目录内容: {files}")
        except Exception as e:
            print(f"读取模型目录失败: {e}")
            return False
        
        # 检查必要的模型文件是否存在（支持多种格式）
        # 检查config文件
        has_config = 'config.json' in files or any('config' in f and f.endswith('.json') for f in files)
        
        # 检查tokenizer文件
        has_tokenizer = ('tokenizer_config.json' in files and 'vocab.json' in files) or \
                      ('tokenizer_config.json' in files and 'vocab.txt' in files) or \
                      any('tokenizer' in f and f.endswith('.json') for f in files)
        
        # 检查模型权重文件（支持多种格式和分片）
        has_model = False
        # 检查safetensors格式
        if any('model-00001-of' in f and f.endswith('.safetensors') for f in files) or \
           any('pytorch_model' in f and f.endswith('.safetensors') for f in files) or \
           'model.safetensors' in files:
            has_model = True
        # 检查bin格式
        elif any('pytorch_model-00001-of' in f and f.endswith('.bin') for f in files) or \
             'pytorch_model.bin' in files:
            has_model = True
        # 检查是否有索引文件和对应的分片文件
        elif 'model.safetensors.index.json' in files and any(f.endswith('.safetensors') for f in files):
            has_model = True
        
        # 更宽松的判断逻辑，确保能识别已下载的模型
        if has_config and has_tokenizer and has_model:
            print(f"✓ 本地模型检查通过")
            return True
        else:
            print(f"✗ 本地模型文件不完整")
            print(f"- 配置文件: {'✓ 存在' if has_config else '✗ 缺失'}")
            print(f"- Tokenizer文件: {'✓ 存在' if has_tokenizer else '✗ 缺失'}")
            print(f"- 模型权重文件: {'✓ 存在' if has_model else '✗ 缺失'}")
            return False
    
    def download_model(self, model_dir=None):
        """
        下载模型到本地，支持魔塔社区下载和本地模型加载
        """
        # 兼容原始配置格式
        if 'model_name' in self.config['model']:
            model_name = self.config['model']['model_name']
        elif 'name' in self.config['model']:
            model_name = self.config['model']['name']
        else:
            model_name = "openbmb/MiniCPM-V-4.5"
            print(f"警告: 配置中未找到模型名称，使用默认值: {model_name}")
        
        # 使用传入的model_dir或尝试从配置中获取
        if model_dir is None:
            if 'model' in self.config and 'local_model_path' in self.config['model']:
                cache_dir = self.config['model']['local_model_path']
            elif 'model' in self.config and 'cache_dir' in self.config['model']:
                cache_dir = self.config['model']['cache_dir']
            else:
                cache_dir = "models"
                print(f"警告: 配置中未找到下载目录，使用默认值: {cache_dir}")
        else:
            cache_dir = model_dir
        
        print(f"准备下载模型: {model_name}")
        print(f"下载目录: {cache_dir}")
        
        try:
            # 创建下载目录
            os.makedirs(cache_dir, exist_ok=True)
            
            # 尝试使用多种下载方式
            download_methods = [
                self._download_from_modelscope,
                self._download_with_git
            ]
            
            success = False
            for method in download_methods:
                try:
                    if method(model_name, cache_dir):
                        success = True
                        break
                except Exception as e:
                    print(f"使用{method.__name__}下载失败: {e}")
                    print("尝试下一种下载方式...")
            
            if not success:
                print("所有下载方式都失败了！请手动下载模型。")
                print(f"推荐从魔塔社区下载: https://modelscope.cn/models")
                print(f"或从Hugging Face下载: https://huggingface.co/{model_name}")
                print(f"下载后请解压到: {cache_dir}")
                return False
            
            print(f"✓ 模型下载完成")
            return True
        except Exception as e:
            print(f"下载模型时出错: {e}")
            return False
    
    def _download_from_modelscope(self, model_name, cache_dir):
        """
        从魔塔社区下载模型
        """
        print("尝试从魔塔社区下载模型...")
        
        try:
            # 检查是否安装了modelscope
            try:
                import modelscope
                from modelscope import snapshot_download
                print("已检测到modelscope库")
            except ImportError:
                print("未安装modelscope库，尝试安装...")
                subprocess.check_call([sys.executable, '-m', 'pip', 'install', 'modelscope'])
                import modelscope
                from modelscope import snapshot_download
                
            # 魔塔社区的模型名称可能与huggingface不同，这里做映射
            modelscope_mapping = {
                "openbmb/MiniCPM-V-4.5": "openbmb/MiniCPM-V-4_5"
            }
            
            # 获取魔塔社区的模型ID
            modelscope_id = modelscope_mapping.get(model_name, model_name.replace('/', '-'))
            print(f"使用魔塔社区模型ID: {modelscope_id}")
            
            # 下载模型，使用适合的参数
            print(f"开始从魔塔社区下载模型...")
            try:
                # 尝试使用常用参数组合
                snapshot_download(
                    model_id=modelscope_id,
                    cache_dir=cache_dir,
                    ignore_file_pattern=["*.md", "*.git*"]
                )
            except Exception as e:
                print(f"第一次尝试下载失败: {e}")
                print("尝试使用另一种参数组合...")
                # 尝试另一种参数组合
                snapshot_download(
                    model_id=modelscope_id,
                    cache_dir=cache_dir
                )
            
            return True
        except Exception as e:
            print(f"魔塔社区下载失败: {e}")
            return False
    
    def _download_with_git(self, model_name, cache_dir):
        """
        使用git下载模型
        """
        print("尝试使用git下载模型...")
        
        # 检查git是否可用
        try:
            subprocess.run(["git", "--version"], check=True, capture_output=True)
        except (subprocess.SubprocessError, FileNotFoundError):
            print("git命令不可用")
            return False
        
        repo_url = f"https://huggingface.co/{model_name}.git"
        temp_dir = os.path.join(os.path.dirname(cache_dir), f"{model_name.split('/')[-1]}_temp")
        
        try:
            # 克隆仓库
            subprocess.run(["git", "clone", repo_url, temp_dir], check=True)
            # 移动文件到目标目录
            os.makedirs(cache_dir, exist_ok=True)
            for item in os.listdir(temp_dir):
                src = os.path.join(temp_dir, item)
                dst = os.path.join(cache_dir, item)
                if os.path.exists(dst):
                    if os.path.isdir(dst):
                        shutil.rmtree(dst)
                    else:
                        os.remove(dst)
                shutil.move(src, cache_dir)
            # 清理临时目录
            shutil.rmtree(temp_dir)
            return True
        except Exception as e:
            print(f"git下载失败: {e}")
            # 清理临时目录
            if os.path.exists(temp_dir):
                shutil.rmtree(temp_dir)
            return False
    
    def _download_with_wget(self, model_name, cache_dir):
        """
        使用wget下载模型（备用方案）
        """
        print("尝试使用wget下载模型（仅作为参考，完整下载需要更多步骤）...")
        print(f"请访问 https://huggingface.co/{model_name}/tree/main 手动下载所需文件")
        print(f"下载后请解压到: {cache_dir}")
        # 这里只是提供信息，实际下载需要更复杂的实现
        return False
    
    def load_model(self, device, dtype):
        """
        加载模型到指定设备，支持模型分片和非分片模型
        
        支持两种模型格式：
        1. 分片模型：由多个文件组成，如model-00001-of-00002.safetensors, model-00002-of-00002.safetensors
        2. 非分片模型：单个model.safetensors文件
        """
        try:
            # 获取模型目录
            if 'model_dir' in self.config['paths']:
                model_dir = self.config['paths']['model_dir']
            else:
                # 兼容原始配置格式
                if 'local_model_path' in self.config['model']:
                    model_dir = self.config['model']['local_model_path']
                else:
                    model_dir = "models"
                    print(f"警告: 配置中未找到模型目录，使用默认值: {model_dir}")
            
            print(f"使用模型目录: {model_dir}")
            
            # 检查模型目录是否存在
            if not os.path.exists(model_dir):
                print(f"错误: 模型目录不存在: {model_dir}")
                print("请检查配置文件中的模型路径是否正确")
                # 尝试创建目录并下载模型
                try:
                    os.makedirs(model_dir, exist_ok=True)
                    print("开始下载模型...")
                    if not self.download_model(model_dir):
                        print("模型下载失败，无法加载模型")
                        return False
                except Exception as e:
                    print(f"创建模型目录失败: {e}")
                    return False
            else:
                # 检查模型类型（分片或非分片）
                is_sharded_model = False
                shard_files = []
                model_files = []
                
                # 检查目录中的文件
                try:
                    files = os.listdir(model_dir)
                    for file in files:
                        file_path = os.path.join(model_dir, file)
                        if os.path.isfile(file_path) and file.endswith(".safetensors"):
                            model_files.append(file)
                            # 检查是否是分片文件，格式如model-00001-of-00002.safetensors
                            if re.search(r"-\d{5}-of-\d{5}\.safetensors$", file):
                                shard_files.append(file)
                except Exception as e:
                    print(f"检查模型文件时出错: {e}")
                
                # 确定模型类型
                if len(shard_files) > 1:
                    is_sharded_model = True
                    print(f"检测到分片模型，包含 {len(shard_files)} 个分片文件")
                elif "model.safetensors" in model_files:
                    print("检测到非分片模型 (model.safetensors)")
                elif len(model_files) > 0:
                    print(f"检测到其他格式模型，包含 {len(model_files)} 个模型文件")
                else:
                    print("警告: 模型目录中未找到safetensors文件")
                
                # 首先尝试直接加载模型
                try:
                    print(f"尝试直接加载本地模型，模型路径: {model_dir}")
                    print(f"设备: {device}, 数据类型: {dtype}")
                    
                    # 尝试加载tokenizer
                    print("加载tokenizer...")
                    self.tokenizer = AutoTokenizer.from_pretrained(
                        model_dir,
                        local_files_only=True,
                        trust_remote_code=True
                    )
                    
                    # 尝试加载模型，根据模型类型调整加载参数
                    print("加载模型...")
                    
                    # 基础参数
                    model_kwargs = {
                        "torch_dtype": dtype,
                        "local_files_only": True,
                        "trust_remote_code": True
                    }
                    
                    # 针对CUDA设备的参数
                    if device.startswith("cuda"):
                        # 对于分片模型，使用更适合的参数
                        if is_sharded_model:
                            model_kwargs.update({
                                "device_map": "auto",
                                "low_cpu_mem_usage": True,
                                "offload_folder": ".offload",  # 为分片模型添加offload支持
                                "offload_state_dict": True
                            })
                        else:
                            # 非分片模型使用标准参数
                            model_kwargs.update({
                                "device_map": "auto",
                                "low_cpu_mem_usage": True
                            })
                        
                        # 针对int4模型的特殊处理
                        if "int4" in model_dir.lower():
                            model_kwargs.update({
                                "load_in_4bit": True,
                                "bnb_4bit_compute_dtype": dtype
                            })
                    else:
                        # 对于CPU设备
                        model_kwargs.update({
                            "device_map": device,
                            "low_cpu_mem_usage": True
                        })
                    
                    # 加载模型
                    self.model = AutoModel.from_pretrained(model_dir, **model_kwargs)
                    
                    # 设置为eval模式
                    self.model.eval()
                    
                    print("✓ 模型加载成功")
                    return True
                except Exception as e:
                    print(f"直接加载模型失败: {e}")
                    print("检查本地模型文件完整性...")
                    
                    # 只有在直接加载失败时，才进行模型文件完整性检查
                    if not self.check_local_model(model_dir):
                        print("本地模型文件不完整，开始下载...")
                        if not self.download_model(model_dir):
                            print("模型下载失败，无法加载模型")
                            return False
                    else:
                        # 如果文件完整性检查通过但仍加载失败，尝试不同的加载参数组合
                        print("本地模型文件完整性检查通过，但加载失败，尝试使用不同的加载参数...")
            
            # 如果前面的尝试都失败了，使用备用加载参数进行最后的尝试
            print(f"开始加载模型，模型路径: {model_dir}")
            print(f"设备: {device}, 数据类型: {dtype}")
            
            # 加载tokenizer
            print("加载tokenizer...")
            self.tokenizer = AutoTokenizer.from_pretrained(
                model_dir,
                local_files_only=True,
                trust_remote_code=True
            )
            
            # 加载模型（使用不同的参数组合）
            print("使用备用参数加载模型...")
            try:
                if device.startswith("cuda"):
                    # 尝试不使用low_cpu_mem_usage参数
                    self.model = AutoModel.from_pretrained(
                        model_dir,
                        torch_dtype=dtype,
                        device_map="auto",
                        local_files_only=True,
                        trust_remote_code=True,
                        load_in_4bit=True if "int4" in model_dir.lower() else False,
                        bnb_4bit_compute_dtype=dtype if "int4" in model_dir.lower() else None
                    )
                else:
                    # 对于CPU设备，直接加载
                    self.model = AutoModel.from_pretrained(
                        model_dir,
                        torch_dtype=dtype,
                        device_map=device,
                        local_files_only=True,
                        trust_remote_code=True
                    )
                
                # 设置为eval模式
                self.model.eval()
                
                print("✓ 模型加载成功")
                return True
            except Exception as e:
                print(f"备用参数加载模型失败: {e}")
                print("尝试使用CPU模式加载，然后手动移至目标设备...")
                
                # 最后的尝试：先在CPU上加载，然后移至目标设备
                try:
                    self.model = AutoModel.from_pretrained(
                        model_dir,
                        torch_dtype=dtype,
                        device_map="cpu",
                        local_files_only=True,
                        trust_remote_code=True
                    )
                    
                    # 手动移至目标设备
                    if device.startswith("cuda"):
                        print(f"将模型从CPU移至{device}...")
                        # 处理meta tensor情况，使用to_empty替代to
                        try:
                            self.model = self.model.to(device)
                        except NotImplementedError:
                            print("检测到meta tensor，使用to_empty方法移动设备...")
                            self.model = self.model.to_empty(device=device)
                    
                    # 设置为eval模式
                    self.model.eval()
                    
                    print("✓ 模型加载成功")
                    return True
                except Exception as e2:
                    print(f"最后的加载尝试也失败了: {e2}")
                    import traceback
                    print(f"错误详情: {traceback.format_exc()}")
                    return False
        except Exception as e:
            print(f"加载模型时出错: {e}")
            import traceback
            print(f"错误详情: {traceback.format_exc()}")
            return False
    
    def get_model(self):
        """
        获取加载的模型
        """
        return self.model
    
    def get_tokenizer(self):
        """
        获取加载的tokenizer
        """
        return self.tokenizer
    
    def unload_model(self):
        """
        卸载模型，释放内存
        """
        if self.model is not None:
            del self.model
            self.model = None
        if self.tokenizer is not None:
            del self.tokenizer
            self.tokenizer = None
        # 清理CUDA缓存
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        print("模型已卸载，内存已释放")